#include "xpeuser.h"
BOOL EnumImportModule(
	PIMAGE_IMPORT_DESCRIPTOR pImportDescriptor,
	char* pRealImageBase,
	ULONG SizeOfImage,
	ULONG_PTR Parameter
	)
{
	char* ModuleName;
	HMODULE hModule;
	PIMAGE_THUNK_DATA pOriginalFirstThunk, pFirstThunk;
	char* EndOfImage;
	char* pProcName;
	BOOL Ret = FALSE;


	EndOfImage = pRealImageBase + SizeOfImage;
	ModuleName = pRealImageBase + pImportDescriptor->Name;
	hModule = LoadLibrary(ModuleName);
	if(!hModule)
		goto Leave;
	pOriginalFirstThunk = (PIMAGE_THUNK_DATA)(pRealImageBase + pImportDescriptor->OriginalFirstThunk);
	pFirstThunk = (PIMAGE_THUNK_DATA)(pRealImageBase + pImportDescriptor->FirstThunk);
	do{
		if((char*)pOriginalFirstThunk > EndOfImage)
			goto Leave;
		if((char*)pFirstThunk > pRealImageBase + SizeOfImage)
			goto Leave;
		if(pOriginalFirstThunk->u1.Function & IMAGE_ORDINAL_FLAG){
			//import by ordinal
			pProcName = (char*)(pOriginalFirstThunk->u1.Ordinal & 0xffff);
		}else{
			//import by function
			pProcName = (char*)(((PIMAGE_IMPORT_BY_NAME)(pRealImageBase + pOriginalFirstThunk->u1.Function))->Name);
		}
		pFirstThunk->u1.Function = (ULONG_PTR)GetProcAddress(hModule, pProcName);
		if(!pFirstThunk->u1.Function)
			goto Leave;
		pOriginalFirstThunk++;
		pFirstThunk++;
	}while(pOriginalFirstThunk->u1.Function);
	Ret = TRUE;
Leave:
	return Ret;
}

char* LoadUserImage(char* pFileName)
{
	HANDLE hFile;
	ULONG FileSize;
	ULONG BytesEffect;
	char* pImageBase;
	char* pRealImageBase = NULL;
	char* pImage;
	PIMAGE_FILE_HEADER pFileHeader;
	PIMAGE_OPTIONAL_HEADER pOptionalHeader;
	ULONG SizeOfImage;
	hFile = CreateFile(pFileName,
		GENERIC_READ,
		FILE_SHARE_READ,
		NULL,
		OPEN_EXISTING,
		0, NULL);
	if(hFile == (HANDLE)-1){
		DbgPrint("Open File Failed\n");
		goto Leave;
	}
	FileSize = GetFileSize(hFile, NULL);
	pImage = new char[FileSize];
	if(!pImage)
		goto NeedCloseFileHandle;
	if(!ReadFile(hFile, pImage, FileSize, &BytesEffect, NULL)){
		DbgPrint("ReadFile Error\n");
		goto NeedFreeFileMem;
	}
	if(!ValidatePeFile(pImage, FileSize)){
		DbgPrint("not a vaild pe file\n");
		goto NeedFreeFileMem;
	}
	pFileHeader = GetFileHeaderFromBase(pImage);
	pOptionalHeader = (PIMAGE_OPTIONAL_HEADER)(pFileHeader + 1);
	SizeOfImage = pOptionalHeader->SizeOfImage;
	pImageBase = (char*)pOptionalHeader->ImageBase;
	pRealImageBase = (char*)VirtualAlloc(pImageBase, SizeOfImage, MEM_RESERVE|MEM_COMMIT, PAGE_EXECUTE_READWRITE);
	if(!pRealImageBase){
		pRealImageBase = (char*)VirtualAlloc(NULL, SizeOfImage, MEM_RESERVE|MEM_COMMIT, PAGE_EXECUTE_READWRITE);
		if(!pRealImageBase)
			goto NeedFreeFileMem;
	}
	if(!LayoutPe(pImage, FileSize, pRealImageBase, SizeOfImage)){
		DbgPrint("Layout PE Failed\n");
		goto NeedFreeModuleMem;
	}
	if(!EnumImportModules(pRealImageBase, EnumImportModule, 0)){
		DbgPrint("EnumImportModules Failed\n");
		goto NeedFreeModuleMem;
	}
	if(!ProcessRelocations(pRealImageBase, pRealImageBase - pImageBase)){
		DbgPrint("LdrProcessRelocations Failed\n");
		goto NeedFreeModuleMem;
	}
	goto NeedFreeFileMem;
NeedFreeModuleMem:
	VirtualFree(pRealImageBase, 0, MEM_RELEASE);
	pRealImageBase = NULL;
NeedFreeFileMem:
	delete []pImage;
NeedCloseFileHandle:
	CloseHandle(hFile);
Leave:
	return pRealImageBase;
}